#!/usr/bin/env python
# coding=utf-8
# @FileName :database.py
# @Time     :2023/4/10 16:47
# @Author   :ZXC
# @Desc     :
from pydantic import BaseModel, Field
from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker

from ..enums.database import DatabaseEnum


class DBObj(BaseModel):
    """
    数据库对象模型
    """
    sql_type: DatabaseEnum = Field(..., description='数据库类型')
    host: str = Field(..., description='数据库服务器的主机名或ip地址')
    port: int = Field(3306, description='数据库服务器上的端口，3306，mysql就是3306，必须是数字')
    username: str = Field(..., description='数据库用户名')
    password: str = Field(..., description='数据库密码')
    database: str = Field(..., description='数据库名称')
    use_ssh: bool = Field(False, description='是否启用SSH隧道堡垒机')
    ssh_host: str = Field('127.0.0.1', description='堡垒机ip地址或主机名')
    ssh_port: int = Field(22, description='堡垒机连接mysql服务器的端口号，一般都是22，必须是数字')
    ssh_username: str = Field('root', description='堡垒机用户名')
    ssh_password: str = Field('123456', description='堡垒机用户密码')

    def get_key(self):
        """
        拼接获取数据库对象模型的key
        :return:
        """
        return f"{self.sql_type.name}-{self.host}-{self.port}-{self.username}-{self.password}-{self.database}-" \
               f"{self.use_ssh}-{self.ssh_host}-{self.ssh_port}-{self.ssh_username}-{self.ssh_password}"


class DBConn(BaseModel):
    """
    数据库连接对象模型
    """
    engine: AsyncEngine = Field(..., description='引擎')
    session: async_sessionmaker = Field(..., description='会话')

    class Config:
        arbitrary_types_allowed = True  # 允许任意类型


class SqlStrObj(BaseModel):
    """
    数据库对象模型
    """
    sql: str = Field(..., description='数据库语句')
